import matplotlib.pyplot as plt
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np
from sklearn import tree

df = pd.read_excel('data2.xlsx')
Ycol = ['y']
Xcols = list(set(df.columns) - set(Ycol))
needToNumCols = ['marital', 'housing', 'pre_outcome', 'contact', 'loan']

def strColToNum(strCol):
    strList = set()
    for i in strCol:
        strList.add(i)
    sub = 0
    print(strList)
    for i in strList:
        strCol = strCol.replace(i, sub)
        sub += 1
    return strCol

for colName in needToNumCols:
    print(colName)
    df[colName] = strColToNum(df[colName])

clf = DecisionTreeClassifier()
Xdata = df[Xcols]
Ydata = df['y']
X_train, X_test, y_train, y_test = train_test_split(Xdata, Ydata, test_size=0.1)
clf.fit(X_train, y_train)
print(clf.score(X_train, y_train))
print(clf.score(X_test, y_test))
# plt.figure(figsize=(30,5))
# tree.plot_tree(clf, fontsize=7, feature_names=Xcols)
# plt.savefig('tree_high_dpi', dpi=100)